import Router from 'koa-router';
import { PassThrough } from 'stream';
import { generate } from '../utils/index.js';
import {getCollection} from '../chroma/index.js'
const knowledge_base_chat = new Router();

knowledge_base_chat.post('/', async(ctx) => {
    ctx.set({
        'Connection': 'keep-alive',
        'Cache-Control': 'no-cache',
        'Content-Type': 'text/event-stream',
    });
    const stream = new PassThrough();
    ctx.body = stream;
    ctx.status = 200;
    const data = ctx.request.body
    const vectorStore = getCollection(data.knowledge_base)
    const docs = await vectorStore.similaritySearchWithScore(data.text,3)
    let prompt = `请严格基于提供的参考文档回答以下问题:“${data.text}”,若参考文档中找不到答案,则回复“根据已知知识无法回答问题”,参考文档为:\n\`\`\`${
        (()=>{
            let tmp = ''
            docs.forEach((doc, index) => { 
                tmp += '参考文档' + String(index + 1) + ':'+ (doc[0].pageContent.replace(/\n/g, " ") + '\n')
            })
            return tmp
        })()
    }\`\`\``
    // console.log(prompt)
    generate(prompt,  [], (e) => {
        stream.write(JSON.stringify(e))
    }, () => {
        stream.end()
    })
})
export default knowledge_base_chat;
